/*
 * Wazuh Vulnerability scanner - Scan Orchestrator
 * Copyright (C) 2015, Wazuh Inc.
 * May 2, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#ifndef _REMEDIATION_DATA_CACHE_HPP
#define _REMEDIATION_DATA_CACHE_HPP

#include "../policyManager/policyManager.hpp"
#include "cacheLRU.hpp"
#include "singleton.hpp"
#include "socketDBWrapper.hpp"
#include "wazuhDBQueryBuilder.hpp"
#include "wdbDataException.hpp"
#include <mutex>
#include <string>
#include <unordered_set>

/**
 * @brief Remediation structure.
 */
struct Remediation final
{
    std::unordered_set<std::string> hotfixes; ///< Installed hotfixes.
};

/**
 * @brief RemediationDataCache class.
 *
 * @note This class queries the Wazuh-DB to get the remediation data for a given agent, and stores it in a LRU cache
 */
template<typename TSocketDBWrapper = SocketDBWrapper>
class RemediationDataCache final : public Singleton<RemediationDataCache<>>
{
private:
    LRUCache<std::string, Remediation> m_remediationData {PolicyManager::instance().getRemediationLRUSize()};
    std::mutex m_mutex;

    Remediation getRemediationDataFromWdb(const std::string& agentId) const
    {
        nlohmann::json response;
        try
        {
            TSocketDBWrapper::instance().query(WazuhDBQueryBuilder::builder().agentGetHotfixesCommand(agentId).build(),
                                               response);
        }
        catch (const SocketDbWrapperException& e)
        {
            throw WdbDataException(e.what(), agentId);
        }
        catch (const std::exception& e)
        {
            throw std::runtime_error("Unable to retrieve remediation data from Wazuh-DB (agent " + agentId +
                                     "). Reason: " + e.what());
        }

        Remediation remediationData;

        if (response.empty())
        {
            // No remediation data found (no hotfixes installed)
            return remediationData;
        }

        // Iterate over the response and store the hotfixes.
        for (auto& hotfix : response)
        {
            if (hotfix.contains("hotfix"))
            {
                remediationData.hotfixes.insert(hotfix.at("hotfix"));
            }
        }

        return remediationData;
    } // LCOV_EXCL_LINE

public:
    /**
     * @brief This method returns the Remediation data.
     * @param agentId agent id.
     *
     * @return Remediation
     */
    Remediation getRemediationData(const std::string& agentId)
    {
        auto getValue = [this](const std::string& agentIdValue) -> std::optional<Remediation>
        {
            std::scoped_lock lock(m_mutex);
            if (auto value = m_remediationData.getValue(agentIdValue); value)
            {
                return *value;
            }
            return std::nullopt;
        };

        if (auto remediationData = getValue(agentId); remediationData != std::nullopt)
        {
            return *remediationData;
        }

        const auto remediationData = getRemediationDataFromWdb(agentId);

        if (!remediationData.hotfixes.empty())
        {
            std::scoped_lock lock(m_mutex);
            // Update the cache with the queried data
            m_remediationData.insertKey(agentId, remediationData);
        }

        return remediationData;
    } // LCOV_EXCL_LINE

    /**
     * @brief Add remediation data to the cache.
     *
     * @note If the agentId already exists in the cache, the new data is merged with the existing data.
     *
     * @param agentId agent id.
     * @param newRemediationData data to be inserted.
     */
    void addRemediationData(const std::string& agentId, Remediation newRemediationData)
    {
        std::scoped_lock lock(m_mutex);

        // Merge the new data with the current data.
        if (auto currentData = m_remediationData.getValue(agentId); currentData.has_value())
        {
            newRemediationData.hotfixes.insert(currentData->hotfixes.begin(), currentData->hotfixes.end());
        }

        m_remediationData.insertKey(agentId, newRemediationData);
    }
};

#endif // _REMEDIATION_DATA_CACHE_HPP
